#include <stdio.h>
#include <stdlib.h>
#include "mpi.h"

/****************************************************
 * \file test_scatterv.c
 * \author James Edmondson
 * \brief Scatters a list of ranks to each rank of
 *        MPI_COMM_WORLD from rank 0. This test differs
 *        from test_scatter.c in that it uses the MPI_Scatterv
 *        function with odd displacements.
 ****************************************************/

int main(int argc, char **argv)
{
	int size, rank, inbuf[2], *outbuf, *displs, *counts, i, success;

	MPI_Init(&argc, &argv);
	
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

	outbuf = (int *) malloc(sizeof(int) * size * 2);
	displs = (int *) malloc(sizeof(int) * size);
	counts = (int *) malloc(sizeof(int) * size);

	success = 1;
	
	/* test 1 */

	for (i = 0; i < size * 2; i++)
	{
		outbuf[i] = i;
	}
	
	for (i = 0; i < size; i++)
	{
		displs[i] = i;
		counts[i] = 1;
	}
	MPI_Scatterv(outbuf, counts, displs, MPI_INT, inbuf, 1, MPI_INT, 0,
		MPI_COMM_WORLD);
	printf("after test 1: inbuf[0] = %d\n", inbuf[0]);
	if (inbuf[0] != rank)
	{
		printf("test 1 failed\n");
		success = 0;
	}

	/* test 2 */

	for (i = 0; i < size; i++)
	{
		displs[i] = i * 2;
		counts[i] = 2;
	}
	MPI_Scatterv(outbuf, counts, displs, MPI_INT, inbuf, 2, MPI_INT, 0,
		MPI_COMM_WORLD);
	printf("after test 2: inbuf[0] = %d inbuf[1] = %d\n", inbuf[0], inbuf[1]);
	if (inbuf[0] != rank * 2 || inbuf[1] != rank * 2 + 1)
	{
		printf("test 2 failed\n");
		success = 0;
	}
	
	/* test 3 */

	for (i = 0; i < size; i++)
	{
		displs[i] = i * 2;
		counts[i] = 2 - (i % 2);
	}
	MPI_Scatterv(outbuf, counts, displs, MPI_INT, inbuf, 2 - (rank % 2), MPI_INT, 0,
		MPI_COMM_WORLD);
	if ((rank % 2) == 0)
	{
		printf("after test 3: inbuf[0] = %d inbuf[1] = %d\n", inbuf[0], inbuf[1]);
		if (inbuf[0] != rank * 2 || inbuf[1] != rank * 2 + 1)
		{
			printf("test 3 failed\n");
			success = 0;
		}
	}
	else
	{
		printf("after test 3: inbuf[0] = %d\n", inbuf[0]);
		if (inbuf[0] != rank * 2)
		{
			printf("test 3 failed\n");
			success = 0;
		}
	}
	
	if (success)
	{
		printf("SUCCESS\n");
	}

	MPI_Finalize();

	return 0;
}

